import gradio as gr
import time
import os
import sys
import threading
import queue
from io import StringIO
from ultralytics import YOLO
import traceback

class OutputCapture:
    """捕获标准输出的类"""
    def __init__(self, textbox_update_func):
        self.textbox_update_func = textbox_update_func
        self.buffer = ""
        self.original_stdout = sys.stdout
        
    def write(self, text):
        self.original_stdout.write(text)
        self.buffer += text
        # 实时更新到gradio textbox
        self.textbox_update_func(self.buffer)
        
    def flush(self):
        self.original_stdout.flush()

def main():
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown("<h1 style='text-align: center;'>YOLOv11 WebUI</h1>")
        with gr.Row():
            with gr.Column(scale=1, min_width=300):
                gr.Markdown("<h2 style='text-align: center;'>任务选择</h2>")
                task_choices = ['detect', 'segment', 'classify', 'pose', 'obb']
                task_select = gr.Radio(
                    choices=task_choices,
                    label="任务",
                    value='detect',
                    interactive=True
                )
            with gr.Column(scale=1, min_width=300):
                gr.Markdown("<h2 style='text-align: center;'>模式选择</h2>")
                mode_choices = ['train', 'val', 'predict', 'export', 'track', 'benchmark']
                mode_select = gr.Radio(
                    choices=mode_choices,
                    label="模式",
                    value='train',
                    interactive=True
                )

        # 训练参数页面
        with gr.Tab("训练参数"):
            with gr.Row():
                with gr.Column(scale=1, min_width=300):
                    gr.Markdown("<h2 style='text-align: center;'>训练参数设置</h2>")
                    with gr.Row():
                        # 获取yolo_models文件夹下所有.pt文件
                        model_dir = "yolo_models"
                        model_files = [f for f in os.listdir(model_dir) if f.endswith('.pt')] if os.path.exists(model_dir) else []
                        model = gr.Dropdown(
                            label="模型(model)", 
                            choices=model_files,
                            value=model_files[0] if model_files else None,
                            interactive=True
                        )
                        # 获取datasets文件夹下所有包含data.yaml的文件夹
                        datasets_dir = "datasets"
                        data_choices = []
                        if os.path.exists(datasets_dir):
                            for folder in os.listdir(datasets_dir):
                                folder_path = os.path.join(datasets_dir, folder)
                                if os.path.isdir(folder_path):
                                    yaml_path = os.path.join(folder_path, "data.yaml")
                                    if os.path.exists(yaml_path):
                                        data_choices.append(f"{folder}/data.yaml")
                        
                        data = gr.Dropdown(
                            label="数据(data)",
                            choices=data_choices,
                            value=data_choices[0] if data_choices else None,
                            interactive=True
                        )
                        epochs = gr.Number(label="训练轮数(epochs)", value=100, precision=0, interactive=True)
                    with gr.Row():
                        time_ = gr.Number(label="最长训练时间小时(time)", value=None, precision=0, interactive=True)
                        patience = gr.Number(label="早停轮数(patience)", value=100, precision=0, interactive=True)
                        batch = gr.Number(label="批大小(batch)", value=8, precision=0, interactive=True)
                    with gr.Row():
                        imgsz = gr.Number(label="输入图片尺寸(imgsz)", value=680, precision=0, interactive=True)
                        with gr.Column():
                            save = gr.Checkbox(label="是否保存检查点(save)", value=True, interactive=True)
                            cache = gr.Checkbox(label="是否启用缓存(cache)", value=False, interactive=True)
                        save_period = gr.Number(label="保存周期(save_period)", value=20, precision=0, interactive=True)
                        device = gr.Textbox(label="设备(device)", value="[0]", interactive=True)
                        workers = gr.Number(label="工作线程数(workers)", value=8, precision=0, interactive=True)
                    with gr.Row():
                        project = gr.Textbox(label="保存训练结果的项目名(project)", value=None, interactive=True)
                        name = gr.Textbox(label="保存训练结果的实验名(name)", value=None, interactive=True)
                        exist_ok = gr.Checkbox(label="是否覆盖已存在的实验(exist_ok)", value=False, interactive=True)
                    with gr.Row():
                        with gr.Column():
                            pretrained = gr.Checkbox(label="是否使用预训练模型(pretrained)", value=True, interactive=True)
                            verbose = gr.Checkbox(label="是否打印详细信息(verbose)", value=True, interactive=True)
                        optimizer = gr.Radio(label="优化器(optimizer)", choices=['SGD', 'Adam', 'Adamax', 'AdamW', 'NAdam', 'RAdam', 'RMSProp', 'auto'], value='auto', interactive=True)
                        seed = gr.Number(label="随机种子(seed)", value=0, precision=0, interactive=True)
                    with gr.Row():
                        with gr.Column():
                            deterministic = gr.Checkbox(label="是否启用确定性模式(deterministic)", value=True, interactive=True)
                            single_cls = gr.Checkbox(label="是否将多类数据训练为单类数据(single_cls)", value=False, interactive=True)
                            rect = gr.Checkbox(label="是否启用矩形训练(rect)", value=False, interactive=True)
                        with gr.Column():
                            cos_lr = gr.Checkbox(label="是否使用余弦学习率(cos_lr)", value=False, interactive=True)
                            resume = gr.Checkbox(label="是否恢复训练(resume)", value=False, interactive=True)
                            amp = gr.Checkbox(label="是否启用自动混合精度(amp)", value=True, interactive=True)
                        with gr.Column():
                            profile = gr.Checkbox(label="是否启用性能分析(profile)", value=False, interactive=True)
                            multi_scale = gr.Checkbox(label="是否启用多尺度训练(multi_scale)", value=False, interactive=True)
                            overlap_mask = gr.Checkbox(label="是否启用重叠掩码(overlap_mask)", value=True, interactive=True)
                        with gr.Column():
                            close_mosaic = gr.Number(label="关闭马赛克(close_mosaic)", value=10, precision=0, interactive=True)
                            fraction = gr.Number(label="数据集比例(fraction)", value=1.0, precision=2, interactive=True)
                    with gr.Row():
                        freeze = gr.Number(label="冻结层数(freeze)", value=None, precision=0, interactive=True)
                        mask_ratio = gr.Number(label="掩码比例(mask_ratio)", value=4, precision=0, interactive=True)
                        dropout = gr.Number(label="dropout", value=0.0, precision=2, interactive=True)

            # 自动传递所有控件的值给回调函数
            train_inputs = [
                model, data, epochs, time_, patience, batch, imgsz, save, cache, save_period, device, workers,
                project, name, exist_ok, pretrained, verbose, optimizer, seed, deterministic, single_cls, rect,
                cos_lr, resume, amp, profile, multi_scale, overlap_mask, close_mosaic, fraction, freeze, mask_ratio, dropout
            ]

            def show_all_params(
                model, data, epochs, time_, patience, batch, imgsz, save, cache, save_period, device, workers,
                project, name, exist_ok, pretrained, verbose, optimizer, seed, deterministic, single_cls, rect,
                cos_lr, resume, amp, profile, multi_scale, overlap_mask, close_mosaic, fraction, freeze, mask_ratio, dropout
            ):
                params = {
                    'model': model,
                    'data': data,
                    'epochs': epochs,
                    'time': time_,
                    'patience': patience,
                    'batch': batch,
                    'imgsz': imgsz,
                    'save': save,
                    'cache': cache,
                    'save_period': save_period,
                    'device': device,
                    'workers': workers,
                    'project': project,
                    'name': name,
                    'exist_ok': exist_ok,
                    'pretrained': pretrained,
                    'verbose': verbose,
                    'optimizer': optimizer,
                    'seed': seed,
                    'deterministic': deterministic,
                    'single_cls': single_cls,
                    'rect': rect,
                    'cos_lr': cos_lr,
                    'resume': resume,
                    'amp': amp,
                    'profile': profile,
                    'multi_scale': multi_scale,
                    'overlap_mask': overlap_mask,
                    'close_mosaic': close_mosaic,
                    'fraction': fraction,
                    'freeze': freeze,
                    'mask_ratio': mask_ratio,
                    'dropout': dropout,
                }
                return str(params)

            result_output = gr.Textbox(label="参数", value=None, interactive=True, lines=3)
            
            # 训练输出文本框
            train_outputs = gr.Textbox(label="训练输出", value="", interactive=True, lines=15, max_lines=20)
            
            def train_with_logging(params_str, progress=gr.Progress()):
                """带实时日志输出的训练函数"""
                try:
                    # 解析参数字符串
                    params_dict = eval(params_str)
                    
                    # 创建YOLO模型实例
                    model_path = os.path.join('yolo_models', params_dict['model']) if not os.path.isabs(params_dict['model']) else params_dict['model']
                    model = YOLO(model_path)
                    
                    # 从params_dict中提取训练参数
                    train_args = {
                        'data': os.path.join('datasets', params_dict['data']),
                        'epochs': params_dict['epochs'],
                        'patience': params_dict['patience'],
                        'batch': params_dict['batch'],
                        'imgsz': params_dict['imgsz'],
                        'save': params_dict['save'],
                        'cache': params_dict['cache'],
                        'device': eval(params_dict['device']) if isinstance(params_dict['device'], str) else params_dict['device'],
                        'workers': params_dict['workers'],
                        'project': params_dict['project'],
                        'name': params_dict['name'],
                        'exist_ok': params_dict['exist_ok'],
                        'pretrained': params_dict['pretrained'],
                        'optimizer': params_dict['optimizer'],
                        'verbose': params_dict['verbose'],
                        'seed': params_dict['seed'],
                        'deterministic': params_dict['deterministic'],
                        'single_cls': params_dict['single_cls'],
                        'rect': params_dict['rect'],
                        'cos_lr': params_dict['cos_lr'],
                        'resume': params_dict['resume'],
                        'amp': params_dict['amp'],
                        'fraction': params_dict['fraction'],
                        'profile': params_dict['profile'],
                        'freeze': params_dict['freeze'],
                        'multi_scale': params_dict['multi_scale'],
                        'overlap_mask': params_dict['overlap_mask'],
                        'mask_ratio': params_dict['mask_ratio'],
                        'dropout': params_dict['dropout'],
                        'close_mosaic': params_dict['close_mosaic']
                    }

                    # 移除None值的参数
                    train_args = {k: v for k, v in train_args.items() if v is not None}
                    
                    # 使用StringIO捕获输出
                    output_buffer = StringIO()
                    original_stdout = sys.stdout
                    
                    try:
                        # 重定向标准输出
                        sys.stdout = output_buffer
                        
                        # 开始训练
                        results = model.train(**train_args)
                        
                        # 获取所有输出
                        output_text = output_buffer.getvalue()
                        
                        return f"训练完成!\n\n{output_text}\n\n最终结果: {results}"
                        
                    finally:
                        # 恢复标准输出
                        sys.stdout = original_stdout
                        
                except Exception as e:
                    return f"训练过程中出现错误: {str(e)}\n\n{traceback.format_exc()}"

            def start_training(params_str):
                """开始训练并实时更新输出"""
                return train_with_logging(params_str)

            show_all_params_btn = gr.Button("显示所有训练参数")
            show_all_params_btn.click(
                fn=show_all_params,
                inputs=train_inputs,
                outputs=result_output
            )
            
            with gr.Row():
                train_btn = gr.Button("开始训练", variant="primary")
                train_btn.click(
                    fn=start_training, 
                    inputs=result_output,
                    outputs=train_outputs
                )
                
                clear_btn = gr.Button("清空输出")
                clear_btn.click(
                    fn=lambda: "",
                    outputs=train_outputs
                )

        # 预测参数页面
        with gr.Tab("预测参数"):
            with gr.Row():
                with gr.Column(scale=1, min_width=300):
                    gr.Markdown("<h2 style='text-align: center;'>预测参数设置</h2>")
                    with gr.Row():
                        pred_model = gr.Textbox(label="模型路径(model)", value=None, interactive=True)
                        pred_source = gr.Textbox(label="预测输入(source)", value=None, interactive=True)
                    with gr.Row():
                        pred_imgsz = gr.Number(label="输入图片尺寸(imgsz)", value=680, precision=0, interactive=True)
                        pred_conf = gr.Number(label="置信度阈值(conf)", value=0.25, precision=2, interactive=True)
                        pred_iou = gr.Number(label="NMS IoU阈值(iou)", value=0.7, precision=2, interactive=True)
                    with gr.Row():
                        pred_device = gr.Textbox(label="设备(device)", value="[0]", interactive=True)
                        pred_project = gr.Textbox(label="保存项目名(project)", value="train_results", interactive=True)
                        pred_name = gr.Textbox(label="保存子目录名(name)", value="predict_output", interactive=True)
                        pred_exist_ok = gr.Checkbox(label="是否覆盖已存在目录(exist_ok)", value=True, interactive=True)
                    with gr.Row():
                        pred_save = gr.Checkbox(label="是否保存结果(save)", value=True, interactive=True)
                        pred_show = gr.Checkbox(label="是否显示结果(show)", value=True, interactive=True)
                        pred_save_txt = gr.Checkbox(label="是否保存txt(save_txt)", value=False, interactive=True)
                        pred_save_conf = gr.Checkbox(label="是否保存置信度(save_conf)", value=False, interactive=True)
                        pred_save_crop = gr.Checkbox(label="是否保存裁剪(save_crop)", value=False, interactive=True)
                        pred_classes = gr.Textbox(label="类别过滤(classes, 逗号分隔)", value=None, interactive=True)
                        pred_half = gr.Checkbox(label="是否使用半精度(half)", value=False, interactive=True)
                        pred_augment = gr.Checkbox(label="是否使用增强(augment)", value=False, interactive=True)
                        pred_visualize = gr.Checkbox(label="是否可视化特征(visualize)", value=False, interactive=True)
                        pred_retina_masks = gr.Checkbox(label="高分辨率掩码(retina_masks)", value=False, interactive=True)
                        pred_line_width = gr.Number(label="框线宽度(line_width)", value=None, precision=0, interactive=True)

            # 自动传递所有预测控件的值给回调函数
            pred_inputs = [
                pred_model, pred_source, pred_imgsz, pred_conf, pred_iou, pred_device, pred_project, pred_name, pred_exist_ok,
                pred_save, pred_show, pred_save_txt, pred_save_conf, pred_save_crop, pred_classes, pred_half, pred_augment,
                pred_visualize, pred_retina_masks, pred_line_width
            ]

            def show_all_pred_params(
                pred_model, pred_source, pred_imgsz, pred_conf, pred_iou, pred_device, pred_project, pred_name, pred_exist_ok,
                pred_save, pred_show, pred_save_txt, pred_save_conf, pred_save_crop, pred_classes, pred_half, pred_augment,
                pred_visualize, pred_retina_masks, pred_line_width
            ):
                params = {
                    'model': pred_model,
                    'source': pred_source,
                    'imgsz': pred_imgsz,
                    'conf': pred_conf,
                    'iou': pred_iou,
                    'device': pred_device,
                    'project': pred_project,
                    'name': pred_name,
                    'exist_ok': pred_exist_ok,
                    'save': pred_save,
                    'show': pred_show,
                    'save_txt': pred_save_txt,
                    'save_conf': pred_save_conf,
                    'save_crop': pred_save_crop,
                    'classes': pred_classes,
                    'half': pred_half,
                    'augment': pred_augment,
                    'visualize': pred_visualize,
                    'retina_masks': pred_retina_masks,
                    'line_width': pred_line_width,
                }
                return str(params)

            pred_result_output = gr.Textbox(label="预测参数", value=None, interactive=True, lines=3)
            show_all_pred_params_btn = gr.Button("显示所有预测参数")
            show_all_pred_params_btn.click(
                fn=show_all_pred_params,
                inputs=pred_inputs,
                outputs=pred_result_output
            )

    demo.launch()

if __name__ == "__main__":
    main()
